import sys
import os
import numpy as np
from sklearn.cluster import KMeans
import cv2
#from utils import *
import time

C, R = 100, 300

def load_image(filename):
    img = cv2.imread(filename)
    img = cv2.resize(img, (C, R))
    return img

def flatten_image(img, refer_mask):
    # 生成消除边缘的mask
    #mask = np.zeros(img.shape, dtype=np.uint8)
    #mult = 4
    #mask[5*mult:47*mult, 1*mult:(C-1*mult+1)] = int(255)

    # 预处理
    # 滤波
    #img = cv2.bilateralFilter(img, 0, 100, 1)
    img = cv2.medianBlur(img, 3)
    #img = cv2.GaussianBlur(img, (5,5), 0)
    #img = cv2.blur(img, (3, 3))

    # 去掉边缘
    img = np.minimum(img, refer_mask)

    # 拉直,供聚类算法使用
    row, col, n = img.shape
    img = img.reshape((row*col, n))

    return img, row, col

def better_image(img, func, refer_mask):
    #mask = np.ones(img.shape, dtype=np.uint8) * 255
    #mult = 4
    #mask[5*mult:47*mult, 1*mult:(C-1*mult+1)] = int(0)
    #img = np.maximum(img, mask)

    mask = np.zeros(img.shape, dtype=np.uint8)
    mult = 4
    mask[5*mult:47*mult, 1*mult:(C-1*mult+1)] = int(255)
    img = np.minimum(img, refer_mask)

    # 预处理
    # 滤波
    #img = cv2.bilateralFilter(img, 0, 100, 1)
    #img = cv2.medianBlur(img, 3)
    # 去掉边缘
    #img = np.minimum(img, mask)

    if func == 'ScaleAbs':
        new_img = cv2.convertScaleAbs(img, alpha=1, beta=100)
    elif func == 'NormalizeL1':
        new_img = cv2.normalize(img, dst=None, norm_type=cv2.NORM_L1, mask=mask)
    elif func == 'NormalizeL2':
        new_img = cv2.normalize(img, dst=None, norm_type=cv2.NORM_L2, mask=mask)
    elif func == 'NormalizeMM':
        new_img = cv2.normalize(img, dst=None, alpha=255, beta=0, norm_type=cv2.NORM_MINMAX, mask=mask)
    elif func == 'EqualizeHist':
        new_img = cv2.equalizeHist(img)
    elif func == 'AdaptHist':
        new_img = cv2.createCLAHE(3,(8,8)).apply(img)
    return new_img

def km_image(img, row, col):
    begin = time.time()
    # 聚类
    km = KMeans(n_clusters=3)
    label = km.fit_predict(img)
    label = label.reshape((row, col))

    #end = time.time()
    #print('km time: {}'.format(end-begin))
    #begin = end

    # 找出灰度值处于中间的类,这是细胞组织部分
    gray = [0, 0, 0]
    gray_cnt = [0, 0, 0]
    find_flag = False
    gray_img = cv2.cvtColor(img.reshape((row, col, 3)), cv2.COLOR_RGB2GRAY)

    #grayImg = better_image(grayImg, 'EqualizeHist')
    for i in range(row):
        for j in range(col):
            idx = label[i, j]
            if gray_cnt[idx] < 5:
                gray[label[i, j]] = (gray_img[i, j] + gray[label[i, j]]) // 2
                gray_cnt[idx] += 1
            if set(gray_cnt) == {5}:
                find_flag = True
                break
        if find_flag:
            break

    mid = sorted(gray)[1]
    midIndex = 2
    for i,v in enumerate(gray):
        if v == mid:
            midIndex = i
            break

    #end = time.time()
    #print('find mid time: {}'.format(end-begin))
    #begin = end

    # 二值化
    output = np.zeros(label.shape, np.uint8)
    for i in range(row):
        for j in range(col):
            if label[i, j] == midIndex:
                output[i, j] = 1
            else:
                output[i, j] = 0

    #end = time.time()
    #print('threshold time: {}'.format(end-begin))
    #begin = end
    return output

def extra_outline(gls, isConvex, num=1):
    # 寻找物体的凸包并绘制凸包的轮廓
    if isConvex:
        calc_num = num
    else:
        calc_num = 1
    for n in range(calc_num):
        contours = cv2.findContours(gls, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        #contours = cv2.findContours(gls, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        if isConvex:
            for cnt in contours[0]:
                hull = cv2.convexHull(cnt)
                length = len(hull)
                # 如果凸包点集中的点个数大于4
                if length > 5:
                    # 绘制图像凸包的轮廓
                    for i in range(length):
                        cv2.line(gls, tuple(hull[i][0]), tuple(hull[(i+1)%length][0]), (255,0,0), 3)
                cv2.drawContours(gls, [cnt], -1, 255, cv2.FILLED)
            if False:
                cv2.imshow(f'gls{n}', gls)

    # 最后再将轮孔中内涂成白色
    contours = cv2.findContours(gls, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(gls, contours[0], -1, 255, cv2.FILLED)
    return gls, contours[0]

def make_mask(img, refer_mask):
    img, row, col = flatten_image(img, refer_mask)

    # 聚类,并二值化
    img = km_image(img, row, col)
    #cv2.imshow(f'km{index}', img*127)

    # 直线检测
    #edges=cv2.Canny(img,1,2,apertureSize=3)
    #minLineLength = 10
    #maxLineGap = 4
    #lines=cv2.HoughLinesP(edges, 1, np.pi/180, 20, minLineLength, maxLineGap)
    ##lines=cv2.HoughLinesP(img, 1, np.pi/180, 5, minLineLength, maxLineGap)
    #if lines is not None:
    #    for l in lines:
    #        x1, y1, x2, y2 = l[0]
    #        print(f'index: {index}: {l[0]}')
    #        cv2.line(img, (x1, y1), (x2, y2), (0, 0, 0), 2)
    #cv2.imshow('edge', edges*128)
    #cv2.imshow(f'km{index}', img*127)

    # 腐蚀
    #kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
    #cv2.imshow('kmimg', img*127)

    area_ratio = len(img.nonzero()[0]) / (row * col)
    #print('area_ratio: {}%'.format(area_ratio))
    if area_ratio > 0.025:
        kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
        mask = cv2.erode(img, kernel)

        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        mask = cv2.erode(mask, kernel)
    elif area_ratio > 0.005:
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (4, 4))
        mask = cv2.erode(img, kernel)
    else:
        kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
        mask = cv2.erode(img, kernel)

    mask = np.clip(mask, 0, 1)

    # 保存对焦区域掩码
    focus_mask = mask
    #cv2.imshow(f'mask{index}', mask*127)

    # 膨胀
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (7,7))
    mask = cv2.dilate(img, kernel)

    # 寻找凸包轮廓, 作为扫描区域
    mask, cnts = extra_outline(mask, True, 1)
    mask = np.clip(mask, 0, 1)

    # 扫描掩码
    scan_mask = mask

    #result = result + mask
    #cv2.imshow(f'result{index}', result*127)
    return focus_mask, scan_mask, img

def get_barcode_width(img, show=False):
    # 假设图片顶部一定有较暗区域
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    min_r, max_r = -1, -1
    isDark = True
    for r in range(gray.shape[0]):
        val = np.median(gray[r])
        if val < 100 and isDark:
            min_r = r
        elif val >= 100:
            max_r = r
            isDark = False
    if min_r >= max_r:
        return -1, -1
    return min_r, max_r

if __name__=='__main__':
    if len(sys.argv)>1:
        imgfile = sys.argv[1]
    else:
        imgfile = 'test.jpg'

    img = load_image(imgfile)

    refer_mask = np.zeros(img.shape, dtype=np.uint8)
    mult = 4
    refer_mask[5*mult:47*mult, 1*mult:(C-1*mult+1)] = int(255)

    focus_mask, scan_mask, raw_mask = make_mask(img, refer_mask)

    min_r, max_r = get_barcode_width(img)

    cv2.imshow(f'img', img)
    cv2.imshow(f'focus_mask', focus_mask * 255)
    cv2.imshow(f'scan_mask', scan_mask * 255)
    cv2.imshow(f'raw_mask', raw_mask * 255)
    cv2.waitKey()
    cv2.destroyAllWindows()

